{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "bccd47fc",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/postgres.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db0855d0",
   "metadata": {},
   "source": [
    "# Postgres Vector Store\n",
    "In this notebook we are going to show how to use [Postgresql](https://www.postgresql.org) and  [pgvector](https://github.com/pgvector/pgvector)  to perform vector searches in LlamaIndex"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e4f33fc9",
   "metadata": {},
   "source": [
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2fc9c18",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-vector-stores-postgres"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "712daea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2d1c538",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import logging\n",
    "# import sys\n",
    "\n",
    "# Uncomment to see debug logs\n",
    "# logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)\n",
    "# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index.core import SimpleDirectoryReader, StorageContext\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.vector_stores.postgres import PGVectorStore\n",
    "import textwrap\n",
    "import openai"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26c71b6d",
   "metadata": {},
   "source": [
    "### Setup OpenAI\n",
    "The first step is to configure the openai key. It will be used to created embeddings for the documents loaded into the index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67b86621",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"<your key>\"\n",
    "openai.api_key = \"<your key>\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "eecf4bd5",
   "metadata": {},
   "source": [
    "Download Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6df9fa89",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p 'data/paul_graham/'\n",
    "!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f7010b1d-d1bb-4f08-9309-a328bb4ea396",
   "metadata": {},
   "source": [
    "### Loading documents\n",
    "Load the documents stored in the `data/paul_graham/` using the SimpleDirectoryReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c154dd4b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Document ID: d05d1211-b9af-4b05-8da6-956e4b389467\n"
     ]
    }
   ],
   "source": [
    "documents = SimpleDirectoryReader(\"./data/paul_graham\").load_data()\n",
    "print(\"Document ID:\", documents[0].doc_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bd24f0a",
   "metadata": {},
   "source": [
    "### Create the Database\n",
    "Using an existing postgres running at localhost, create the database we'll be using."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d61e73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import psycopg2\n",
    "\n",
    "connection_string = \"postgresql://postgres:password@localhost:5432\"\n",
    "db_name = \"vector_db\"\n",
    "conn = psycopg2.connect(connection_string)\n",
    "conn.autocommit = True\n",
    "\n",
    "with conn.cursor() as c:\n",
    "    c.execute(f\"DROP DATABASE IF EXISTS {db_name}\")\n",
    "    c.execute(f\"CREATE DATABASE {db_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0232fd1",
   "metadata": {},
   "source": [
    "### Create the index\n",
    "Here we create an index backed by Postgres using the documents loaded previously. PGVectorStore takes a few arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8731da62",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1358441c0c864d87a860820ed8cf2b2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Parsing documents into nodes:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "69a77c77ff1c48cc8107b445ae4fa0cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating embeddings:   0%|          | 0/17 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from sqlalchemy import make_url\n",
    "\n",
    "url = make_url(connection_string)\n",
    "vector_store = PGVectorStore.from_params(\n",
    "    database=db_name,\n",
    "    host=url.host,\n",
    "    password=url.password,\n",
    "    port=url.port,\n",
    "    user=url.username,\n",
    "    table_name=\"paul_graham_essay\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context, show_progress=True\n",
    ")\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ee4473a-094f-4d0a-a825-e1213db07240",
   "metadata": {},
   "source": [
    "### Query the index\n",
    "We can now ask questions using our index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a2bcc07",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What did the author do?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cf55bf7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author worked on writing and programming before college. They wrote short stories and tried\n",
      "writing programs on an IBM 1401 computer. They also built a microcomputer and started programming on\n",
      "it, writing simple games and a word processor. In college, the author initially planned to study\n",
      "philosophy but switched to AI. They were inspired by a novel called The Moon is a Harsh Mistress and\n",
      "a PBS documentary featuring Terry Winograd using SHRDLU.\n"
     ]
    }
   ],
   "source": [
    "print(textwrap.fill(str(response), 100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68cbd239-880e-41a3-98d8-dbb3fab55431",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What happened in the mid 1980s?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdf5287f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "In the mid-1980s, the author spent a significant amount of time working on a book called \"On Lisp\"\n",
      "and had obtained a contract to publish it. They were paid large amounts of money for their work,\n",
      "which allowed them to save enough to go back to RISD (Rhode Island School of Design) and pay off\n",
      "their college loans. They also learned valuable lessons during this time, such as the importance of\n",
      "having technology companies run by product people rather than sales people, the drawbacks of editing\n",
      "code by too many people, and the significance of being the \"entry level\" option in a competitive\n",
      "market.\n"
     ]
    }
   ],
   "source": [
    "print(textwrap.fill(str(response), 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3bed9e1",
   "metadata": {},
   "source": [
    "### Querying existing index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b2634b",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store = PGVectorStore.from_params(\n",
    "    database=\"vector_db\",\n",
    "    host=\"localhost\",\n",
    "    password=\"password\",\n",
    "    port=5432,\n",
    "    user=\"postgres\",\n",
    "    table_name=\"paul_graham_essay\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    ")\n",
    "\n",
    "index = VectorStoreIndex.from_vector_store(vector_store=vector_store)\n",
    "query_engine = index.as_query_engine()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7075af3-156e-4bde-8f76-6d9dee86861f",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What did the author do?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b088c090",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The author worked on writing and programming before college. They wrote short stories and tried\n",
      "writing programs on an IBM 1401 computer. They also built a microcomputer and started programming on\n",
      "it, writing simple games and a word processor. In college, the author initially planned to study\n",
      "philosophy but switched to AI due to their interest in intelligent computers. They taught themselves\n",
      "AI by learning Lisp.\n"
     ]
    }
   ],
   "source": [
    "print(textwrap.fill(str(response), 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55745895-8f01-4275-abaa-b2ebef2cb4c7",
   "metadata": {},
   "source": [
    "### Hybrid Search  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91cae40f-3cd4-4403-8af4-aca2705e96a2",
   "metadata": {},
   "source": [
    "To enable hybrid search, you need to:\n",
    "1. pass in `hybrid_search=True` when constructing the `PGVectorStore` (and optionally configure `text_search_config` with the desired language)\n",
    "2. pass in `vector_store_query_mode=\"hybrid\"` when constructing the query engine (this config is passed to the retriever under the hood). You can also optionally set the `sparse_top_k` to configure how many results we should obtain from sparse text search (default is using the same value as `similarity_top_k`). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a7e133-39da-40c5-b2c5-7af2c0a3a792",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/suo/dev/llama_index/llama_index/vector_stores/postgres.py:217: SAWarning: TypeDecorator TSVector() will not produce a cache key because the ``cache_ok`` attribute is not set to True.  This can have significant performance implications including some performance degradations in comparison to prior SQLAlchemy versions.  Set this attribute to True if this type object's state is safe to use in a cache key, or False to disable this warning. (Background on this warning at: https://sqlalche.me/e/20/cprf)\n",
      "  session.commit()\n"
     ]
    }
   ],
   "source": [
    "from sqlalchemy import make_url\n",
    "\n",
    "url = make_url(connection_string)\n",
    "hybrid_vector_store = PGVectorStore.from_params(\n",
    "    database=db_name,\n",
    "    host=url.host,\n",
    "    password=url.password,\n",
    "    port=url.port,\n",
    "    user=url.username,\n",
    "    table_name=\"paul_graham_essay_hybrid_search\",\n",
    "    embed_dim=1536,  # openai embedding dimension\n",
    "    hybrid_search=True,\n",
    "    text_search_config=\"english\",\n",
    ")\n",
    "\n",
    "storage_context = StorageContext.from_defaults(\n",
    "    vector_store=hybrid_vector_store\n",
    ")\n",
    "hybrid_index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f8edee4-6c19-4d99-b602-110bdc5708e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "hybrid_query_engine = hybrid_index.as_query_engine(\n",
    "    vector_store_query_mode=\"hybrid\", sparse_top_k=2\n",
    ")\n",
    "hybrid_response = hybrid_query_engine.query(\n",
    "    \"Who does Paul Graham think of with the word schtick\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd454b25-b66c-4733-8ff4-24fb2ee84cec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Roy Lichtenstein\n"
     ]
    }
   ],
   "source": [
    "print(hybrid_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b274ecb",
   "metadata": {},
   "source": [
    "### PgVector Query Options"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a490a0fa",
   "metadata": {},
   "source": [
    "#### IVFFlat Probes\n",
    "\n",
    "Specify the number of [IVFFlat probes](https://github.com/pgvector/pgvector?tab=readme-ov-file#query-options) (1 by default)\n",
    "\n",
    "When retrieving from the index, you can specify an appropriate number of IVFFlat probes (higher is better for recall, lower is better for speed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "111a3682",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(\n",
    "    vector_store_query_mode=query_mode,\n",
    "    similarity_top_k=top_k,\n",
    "    vector_store_kwargs={\"ivfflat_probes\": 10},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6104ef8d",
   "metadata": {},
   "source": [
    "#### HNSW EF Search\n",
    "\n",
    "Specify the size of the dynamic [candidate list](https://github.com/pgvector/pgvector?tab=readme-ov-file#query-options-1) for search (40 by default)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3a44758",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = index.as_retriever(\n",
    "    vector_store_query_mode=query_mode,\n",
    "    similarity_top_k=top_k,\n",
    "    vector_store_kwargs={\"hnsw_ef_search\": 300},\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
